{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/huggingface/transformers/blob/update_notebook/notebooks/05_benchmark.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jG-SjOQTskcX"
   },
   "source": [
    "## **How to benchmark models with Transformers**\n",
    "\n",
    "With ever-larger language models, it is no longer enough to just \n",
    "compare models on their performance on a specific task. One should always be aware of the computational cost that is attached to a specific model. For a given computation environment (*e.g.* type of GPU), the computational cost of training a model or deploying it in inference usually depends only on **the required memory** and **the required time**. \n",
    "\n",
    "Being able to accurately benchmark language models on both *speed* and *required memory* is therefore very important.\n",
    "\n",
    "HuggingFace's Transformer library allows users to benchmark models for both TensorFlow 2 and PyTorch using the `PyTorchBenchmark` and `TensorFlowBenchmark` classes.\n",
    "\n",
    "The currently available features for `PyTorchBenchmark` are summarized in the following table.\n",
    "\n",
    "\n",
    "| | CPU | CPU + torchscript | GPU | GPU + torchscript | GPU + FP16 | TPU |\n",
    ":-- | :--- | :--- | :--- | :--- | :--- | :--- |\n",
    "**Speed - Inference** | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ |\n",
    "**Memory - Inference** | ✔ | ✔ | ✔ | ✔ | ✔ | ✘ |\n",
    "**Speed - Train** | ✔ | ✘ | ✔ | ✘ | ✔ | ✔ |\n",
    "**Memory - Train** | ✔ | ✘ | ✔ | ✘ | ✔ | ✘ |\n",
    "\n",
    "\n",
    "*   *FP16* stands for mixed-precision meaning that computations within the model are done using a mixture of 16-bit and 32-bit floating-point operations, see [here](https://pytorch.org/docs/stable/nn.html#torch.nn.Module.half) for more detail.\n",
    "\n",
    "*   *torchscript* corresponds to PyTorch's torchscript format, see [here](https://pytorch.org/docs/stable/jit.html).\n",
    "\n",
    "The currently available features for `TensorFlowBenchmark` are summarized in the following table.\n",
    "\n",
    "| | CPU | CPU + eager execution | GPU | GPU + eager execution | GPU + XLA | GPU + FP16 | TPU |\n",
    ":-- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n",
    "**Speed - Inference** | ✔ | ✔ | ✔ | ✔ | ✔ | ✘ | ✔ |\n",
    "**Memory - Inference** | ✔ | ✔ | ✔ | ✔ | ✔ | ✘ | ✘ |\n",
    "**Speed - Train** | ✔ | ✘ | ✔ | ✘ | ✘ | ✘ | ✔ |\n",
    "**Memory - Train** | ✔ | ✘ | ✔ | ✘ | ✘ | ✘ | ✘ |\n",
    "\n",
    "*   *eager execution* means that the function is run in the eager execution environment of TensorFlow 2, see [here](https://www.tensorflow.org/guide/eager).\n",
    "\n",
    "*   *XLA* stands for TensorFlow's Accelerated Linear Algebra (XLA) compiler, see [here](https://www.tensorflow.org/xla)\n",
    "\n",
    "*   *FP16* stands for TensorFlow's mixed-precision package and is analogous to PyTorch's FP16 feature, see [here](https://www.tensorflow.org/guide/mixed_precision).\n",
    "\n",
    "***Note***: Benchmark training in TensorFlow is not included in v3.0.2, but available in master.\n",
    "\n",
    "\n",
    "This notebook will show the user how to use `PyTorchBenchmark` and `TensorFlowBenchmark` for two different scenarios:\n",
    "\n",
    "1. **Inference - Pre-trained Model Comparison** - *A user wants to implement a pre-trained model in production for inference. She wants to compare different models on speed and required memory.*\n",
    "\n",
    "2. **Training - Configuration Comparison** - *A user wants to train a specific model and searches that for himself most effective model configuration.*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j-jvAvZ1-GIh"
   },
   "source": [
    "### **Inference - Pre-trained Model Comparison**\n",
    "\n",
    "Let's say we want to employ a question-answering model in production. The questions are expected to be of the same format as in **SQuAD v2**, so that the model to choose should have been fine-tuned on this dataset. \n",
    "\n",
    "HuggingFace's new dataset [webpage](https://huggingface.co/datasets) lets the user see all relevant information about a dataset and even links the models that have been fine-tuned on this specific dataset. Let's check out the dataset webpage of SQuAD v2 [here](https://huggingface.co/datasets/squad_v2).\n",
    "\n",
    "Nice, we can see that there are 7 available models.\n",
    "\n",
    "![Texte alternatif…](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/squad_v2_dataset.png)\n",
    "\n",
    "Let's assume that we have decided to restrict our pipeline to \"encoder-only\" models so that we are left with:\n",
    "\n",
    "- `a-ware/roberta-large-squad-classification`\n",
    "- `a-ware/xlmroberta-squadv2`\n",
    "- `aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2`\n",
    "- `deepset/roberta-base-squad2`\n",
    "- `mrm8488/longformer-base-4096-finetuned-squadv2`\n",
    "\n",
    "Great! In this notebook, we will now benchmark these models on both peak memory consumption and inference time to decide which model should be employed in production.\n",
    "\n",
    "***Note***: None of the models has been tested on performance so that we will just assume that all models perform more or less equally well. The purpose of this notebook is not to find the best model for SQuAD v2, but to showcase how Transformers benchmarking tools can be leveraged.\n",
    "\n",
    "First, we assume to be limited by the available GPU on this google colab, which in this copy amounts to 16 GB of RAM."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2l9C7d7K5-G4"
   },
   "source": [
    "In a first step, we will check which models are the most memory-efficient ones.\n",
    "Let's make sure 100% of the GPU is available to us in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 67
    },
    "colab_type": "code",
    "id": "M7cQmgM5TvlO",
    "outputId": "2797c14e-a62d-42cc-97a6-6c61b015d569"
   },
   "outputs": [],
   "source": [
    "#@title Check available memory of GPU\n",
    "# Check that we are using 100% of GPU\n",
    "# memory footprint support libraries/code\n",
    "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
    "!pip -q install gputil\n",
    "!pip -q install psutil\n",
    "!pip -q install humanize\n",
    "import psutil\n",
    "import humanize\n",
    "import os\n",
    "import GPUtil as GPU\n",
    "GPUs = GPU.getGPUs()\n",
    "# XXX: only one GPU on Colab and isn’t guaranteed\n",
    "gpu = GPUs[0]\n",
    "def printm():\n",
    " process = psutil.Process(os.getpid())\n",
    " print(\"Gen RAM Free: \" + humanize.naturalsize( psutil.virtual_memory().available ), \" | Proc size: \" + humanize.naturalsize( process.memory_info().rss))\n",
    " print(\"GPU RAM Free: {0:.0f}MB | Used: {1:.0f}MB | Util {2:3.0f}% | Total {3:.0f}MB\".format(gpu.memoryFree, gpu.memoryUsed, gpu.memoryUtil*100, gpu.memoryTotal))\n",
    "printm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NuS2CKuQ4qSk"
   },
   "outputs": [],
   "source": [
    "# If GPU RAM Util > 0% => crash notebook on purpose\n",
    "# !kill -9 -1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ikdYDXsj6Nzv"
   },
   "source": [
    "Looks good! Now we import `transformers` and download the scripts `run_benchmark.py`, `run_benchmark_tf.py`, and `plot_csv_file.py` which can be found under `transformers/examples/benchmarking`.\n",
    "\n",
    "`run_benchmark_tf.py` and `run_benchmark.py` are very simple scripts leveraging the `PyTorchBenchmark` and `TensorFlowBenchmark` classes, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "Dylftiyd1IG1"
   },
   "outputs": [],
   "source": [
    "# install transformes\n",
    "!pip uninstall -y transformers\n",
    "!pip install -q git+https://github.com/huggingface/transformers.git\n",
    "\n",
    "# install py3nvml to track GPU memory usage\n",
    "!pip install -q py3nvml\n",
    "\n",
    "!rm -f run_benchmark.py\n",
    "!rm -f run_benchmark_tf.py\n",
    "!rm -f plot_csv_file.py\n",
    "!wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/benchmarking/run_benchmark.py -qq\n",
    "!wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/benchmarking/run_benchmark_tf.py -qq\n",
    "!wget https://raw.githubusercontent.com/huggingface/transformers/master/examples/benchmarking/plot_csv_file.py -qq\n",
    "\n",
    "# import pandas to pretty print csv files\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "C4nz5nGFkOrK"
   },
   "source": [
    "Information about the input arguments to the *run_benchmark* scripts can be accessed by running `!python run_benchmark.py --help` for PyTorch and `!python run_benchmark_tf.py --help` for TensorFlow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "zu7Oufe0jcAj",
    "outputId": "bc52dea5-b721-410c-cf3b-8a7b983a558e"
   },
   "outputs": [],
   "source": [
    "!python run_benchmark.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Q_3TZshjcrjP"
   },
   "source": [
    "Great, we are ready to run our first memory benchmark. By default, both the *required memory* and *time* for inference is enabled. To disable benchmarking on *time*, we add `--no_speed`.\n",
    "\n",
    "The only required parameter is `--models` which expects a list of model identifiers as defined on the [model hub](https://huggingface.co/models). Here we add the five model identifiers listed above.\n",
    "\n",
    "Next, we define the `sequence_lengths` and `batch_sizes` for which the peak memory is calculated.\n",
    "\n",
    "Finally, because the results should be stored in a *CSV* file, the option `--save_to_csv` is added and the path to save the results is added via the `--inference_memory_csv_file` argument. \n",
    "Whenever a benchmark is run, the environment information, *e.g.* GPU type, library versions, ... can be saved using the `--env_info_csv_file` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ykJqt7MEbHIq"
   },
   "outputs": [],
   "source": [
    "# create plots folder in content\n",
    "!mkdir -p plots_pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TSJgpQxBe-Fj"
   },
   "outputs": [],
   "source": [
    "# run benchmark\n",
    "!python run_benchmark.py --no_speed --save_to_csv \\\n",
    "                                --models a-ware/roberta-large-squad-classification \\\n",
    "                                  a-ware/xlmroberta-squadv2 \\\n",
    "                                  aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2 \\\n",
    "                                  deepset/roberta-base-squad2 \\\n",
    "                                  mrm8488/longformer-base-4096-finetuned-squadv2 \\\n",
    "                                --sequence_lengths 32 128 512 1024 \\\n",
    "                                --batch_sizes 32 \\\n",
    "                                --inference_memory_csv_file plots_pt/required_memory.csv \\\n",
    "                                --env_info_csv_file plots_pt/env.csv >/dev/null 2>&1  # redirect all prints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ESHrlnKik396"
   },
   "source": [
    "Under `plots_pt`, two files are now created: `required_memory.csv` and `env.csv`. Let's check out `required_memory.csv` first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 639
    },
    "colab_type": "code",
    "id": "rPg_7fPnuDUa",
    "outputId": "b6272763-7235-43c6-c457-0a4a13bb02e5"
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('plots_pt/required_memory.csv')\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "o2LnaVpyW9TB"
   },
   "source": [
    "Each row in the csv file lists one data point showing the *peak memory* usage for a given model, batch_size and sequence_length. As can be seen, some values have a *NaN* result meaning that an *Out-of-Memory* Error occurred. To better visualize the results, one can make use of the `plot_csv_file.py` script.\n",
    "\n",
    "Before, let's take a look at the information about our computation environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 639
    },
    "colab_type": "code",
    "id": "y6n49pbIXI6E",
    "outputId": "495f011c-87c9-43a1-e1d4-a6501c327e76"
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('plots_pt/env.csv')\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "z316Xf2oXTZz"
   },
   "source": [
    "We can see all relevant information here: the PyTorch version, the Python version, the system, the type of GPU, and available RAM on the GPU, etc...\n",
    "\n",
    "**Note**: A different GPU is likely assigned to a copy of this notebook, so that all of the following results may be different. It is very important to always include the environment information when benchmarking your models for both reproducibility and transparency to other users.\n",
    "\n",
    "Alright, let's plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 534
    },
    "colab_type": "code",
    "id": "yHYUqRzWy8sp",
    "outputId": "22499f33-bafc-42b3-f1b7-fcb202df9cd2"
   },
   "outputs": [],
   "source": [
    "# plot graph and save as image\n",
    "!python plot_csv_file.py --csv_file plots_pt/required_memory.csv --figure_png_file=plots_pt/required_memory_plot.png --no_log_scale --short_model_names a-ware-roberta a-aware-xlm aodiniz-bert deepset-roberta mrm8488-long\n",
    "\n",
    "# show image\n",
    "from IPython.display import Image\n",
    "Image('plots_pt/required_memory_plot.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RKZhRMmJmNH_"
   },
   "source": [
    "At this point, it is important to understand how the peak memory is measured. The benchmarking tools measure the peak memory usage the same way the command `nvidia-smi` does - see [here](https://developer.nvidia.com/nvidia-system-management-interface) for more information. \n",
    "In short, all memory that is allocated for a given *model identifier*, *batch size* and *sequence length* is measured in a separate process. This way it can be ensured that there is no previously unreleased memory falsely included in the measurement. One should also note that the measured memory even includes the memory allocated by the CUDA driver to load PyTorch and TensorFlow and is, therefore, higher than library-specific memory measurement function, *e.g.* this one for [PyTorch](https://pytorch.org/docs/stable/cuda.html#torch.cuda.max_memory_allocated).\n",
    "\n",
    "Alright, let's analyze the results. It can be noted that the models `aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2` and `deepset/roberta-base-squad2` require significantly less memory than the other three models. Besides `mrm8488/longformer-base-4096-finetuned-squadv2` all models more or less follow the same memory consumption pattern with `aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2` seemingly being able to better scale to larger sequence lengths. \n",
    "`mrm8488/longformer-base-4096-finetuned-squadv2` is a *Longformer* model, which makes use of *LocalAttention* (check [this](https://huggingface.co/blog/reformer) blog post to learn more about local attention) so that the model scales much better to longer input sequences.\n",
    "\n",
    "For the sake of this notebook, we assume that the longest required input will be less than 512 tokens so that we settle on the models `aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2` and `deepset/roberta-base-squad2`. \n",
    "\n",
    "To better understand how many API requests of our *question-answering* pipeline can be run in parallel, we are interested in finding out how many batches the two models run out of memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 356
    },
    "colab_type": "code",
    "id": "9Nwmb57M4wIG",
    "outputId": "4c074607-5200-4cca-bbd5-c39d32ce0451"
   },
   "outputs": [],
   "source": [
    "!python run_benchmark.py --no_speed --save_to_csv \\\n",
    "                                --inference_memory_csv_file plots_pt/required_memory_2.csv \\\n",
    "                                --env_info_csv_file plots_pt/env.csv \\\n",
    "                                --models aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2 \\\n",
    "                                  deepset/roberta-base-squad2 \\\n",
    "                                --sequence_lengths 512 \\\n",
    "                                --batch_sizes 64 128 256 512\\\n",
    "                                --no_env_print"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P4JFKLZXqmss"
   },
   "source": [
    "Let's plot the results again, this time changing the x-axis to `batch_size` however."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 534
    },
    "colab_type": "code",
    "id": "tNtvHpE67pgH",
    "outputId": "092c4dac-5002-4603-8eba-cd4bca727744"
   },
   "outputs": [],
   "source": [
    "# plot graph and save as image\n",
    "!python plot_csv_file.py --csv_file plots_pt/required_memory_2.csv \\\n",
    "                          --figure_png_file=plots_pt/required_memory_plot_2.png \\\n",
    "                          --no_log_scale \\\n",
    "                          --short_model_names aodiniz-bert deepset-roberta \\\n",
    "                          --plot_along_batch\n",
    "\n",
    "# show image\n",
    "from IPython.display import Image\n",
    "Image('plots_pt/required_memory_plot_2.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bdoTRF7Yq8oV"
   },
   "source": [
    "Interesting! `aodiniz/bert_uncased_L-10_H-51` clearly scales better for higher batch sizes and does not even run out of memory for 512 tokens.\n",
    "\n",
    "For comparison, let's run the same benchmarking on TensorFlow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 726
    },
    "colab_type": "code",
    "id": "752y4onm-gpy",
    "outputId": "a65c4bc1-f88e-46ae-cb80-27e29a0a1954"
   },
   "outputs": [],
   "source": [
    "# create plots folder in content\n",
    "!mkdir -p plots_tf\n",
    "\n",
    "!TF_CPP_MIN_LOG_LEVEL=3 python run_benchmark_tf.py --no_speed --save_to_csv \\\n",
    "                                --inference_memory_csv_file plots_tf/required_memory_2.csv \\\n",
    "                                --env_info_csv_file plots_tf/env.csv \\\n",
    "                                --models aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2 \\\n",
    "                                         deepset/roberta-base-squad2 \\\n",
    "                                --sequence_lengths 512 \\\n",
    "                                --batch_sizes 64 128 256 512 \\\n",
    "                                --no_env_print \\"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3h5JqW2osAQ7"
   },
   "source": [
    "Let's see the same plot for TensorFlow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 534
    },
    "colab_type": "code",
    "id": "hkw-EOOvA52R",
    "outputId": "3947ccf0-b91c-43bf-8569-d6afe0232185"
   },
   "outputs": [],
   "source": [
    "# plot graph and save as image\n",
    "!python plot_csv_file.py --csv_file plots_tf/required_memory_2.csv --figure_png_file=plots_tf/required_memory_plot_2.png --no_log_scale --short_model_names aodiniz-bert deepset-roberta --plot_along_batch\n",
    "\n",
    "# show image\n",
    "from IPython.display import Image\n",
    "Image('plots_tf/required_memory_plot_2.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ybqol62LsVrF"
   },
   "source": [
    "The model implemented in TensorFlow requires more memory than the one implemented in PyTorch. Let's say for whatever reason we have decided to use TensorFlow instead of PyTorch. \n",
    "\n",
    "The next step is to measure the inference time of these two models. Instead of disabling time measurement with `--no_speed`, we will now disable memory measurement with `--no_memory`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 302
    },
    "colab_type": "code",
    "id": "m8qfllt9uPZg",
    "outputId": "b185f547-fbe6-4287-b8a0-6229d3eec377"
   },
   "outputs": [],
   "source": [
    "!TF_CPP_MIN_LOG_LEVEL=3 python run_benchmark_tf.py --no_memory --save_to_csv \\\n",
    "                                --inference_time_csv_file plots_tf/time_2.csv \\\n",
    "                                --env_info_csv_file plots_tf/env.csv \\\n",
    "                                --models aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2 \\\n",
    "                                         deepset/roberta-base-squad2 \\\n",
    "                                --sequence_lengths 8 32 128 512 \\\n",
    "                                --batch_sizes 256 \\\n",
    "                                --no_env_print \\"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 534
    },
    "colab_type": "code",
    "id": "-bPClv873lrW",
    "outputId": "152f14c7-288a-4471-9cc0-5108cb24804c"
   },
   "outputs": [],
   "source": [
    "# plot graph and save as image\n",
    "!python plot_csv_file.py --csv_file plots_tf/time_2.csv --figure_png_file=plots_tf/time_plot_2.png --no_log_scale --short_model_names aodiniz-bert deepset-roberta --is_time\n",
    "\n",
    "# show image\n",
    "from IPython.display import Image\n",
    "Image('plots_tf/time_plot_2.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "f9sIjRWd4Me1"
   },
   "source": [
    "Ok, this took some time... time measurements take much longer than memory measurements because the forward pass is called multiple times for stable results. Timing measurements leverage Python's [timeit module](https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat) and run 10 times the value given to the `--repeat` argument (defaults to 3), so in our case 30 times.\n",
    "\n",
    "Let's focus on the resulting plot. It becomes obvious that `aodiniz/bert_uncased_L-10_H-51` is around twice as fast as `deepset/roberta-base-squad2`. Given that the model is also more memory efficient and assuming that the model performs reasonably well, for the sake of this notebook we will settle on `aodiniz/bert_uncased_L-10_H-51`. Our model should be able to process input sequences of up to 512 tokens. Latency time of around 2 seconds might be too long though, so let's compare the time for different batch sizes and using TensorFlows XLA package for more speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 202
    },
    "colab_type": "code",
    "id": "aPeMsHJb3t2g",
    "outputId": "56276801-6d56-444c-8ac8-75471136aa84"
   },
   "outputs": [],
   "source": [
    "!TF_CPP_MIN_LOG_LEVEL=3 python run_benchmark_tf.py --no_memory --save_to_csv \\\n",
    "                                --inference_time_csv_file plots_tf/time_xla_1.csv \\\n",
    "                                --env_info_csv_file plots_tf/env.csv \\\n",
    "                                --models aodiniz/bert_uncased_L-10_H-512_A-8_cord19-200616_squad2 \\\n",
    "                                --sequence_lengths 512 \\\n",
    "                                --batch_sizes 8 64 256 \\\n",
    "                                --no_env_print \\\n",
    "                                --use_xla"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_KrzL6y_6Z2T"
   },
   "source": [
    "First of all, it can be noted that XLA reduces latency time by a factor of ca. 1.3 (which is more than observed for other models by TensorFlow [here](https://www.tensorflow.org/xla)). A batch size of 64 looks like a good choice. More or less half a second for the forward pass is good enough.\n",
    "\n",
    "Cool, now it should be straightforward to benchmark your favorite models. All the inference time measurements can also be done using the `run_benchmark.py` script for PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Drht35ylINuK"
   },
   "source": [
    "### **Training - Configuration Comparison**\n",
    "\n",
    "Next, we will look at how a model can be benchmarked on different configurations. This is especially helpful when one wants to decide how to most efficiently choose the model's configuration parameters for training.\n",
    "In the following different configurations of a *Bart MNLI* model will be compared to each other using `PyTorchBenchmark`. \n",
    "\n",
    "Training in `PyTorchBenchmark` is defined by running one forward pass to compute the loss: `loss = model(input_ids, labels=labels)[0]` and one backward pass to compute the gradients `loss.backward()`.\n",
    "\n",
    "Let's see how to most efficiently train a Bart MNLI model from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YTKW0Ml3Wpwq"
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "from transformers import BartConfig, PyTorchBenchmark, PyTorchBenchmarkArguments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6Uw92tMRq6MV"
   },
   "source": [
    "For the sake of the notebook, we assume that we are looking for a more efficient version of Facebook's `bart-large-mnli` model.\n",
    "Let's load its configuration and check out the important parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 637,
     "referenced_widgets": [
      "975f42d7b55c4d0caf229cd4c16df5d2",
      "69b36685703342eaa80b6f0e01f94e04",
      "c8acb33d6a254607a6340c0aa33446f3",
      "a6c3647736554beea36db798827203b2",
      "e812aaf8214c4ad983f41804cb82562b",
      "eed2ce14188a453ca296601ca39133b6",
      "548f91729b8d4f3aa81f78c7a1620101",
      "900c1cb473f54b48a59226c61fafd626"
     ]
    },
    "colab_type": "code",
    "id": "nukyLU7iXBzN",
    "outputId": "ae4ecae5-bd30-4eb4-e4b3-34447036e98d"
   },
   "outputs": [],
   "source": [
    "BartConfig.from_pretrained(\"facebook/bart-large-mnli\").to_diff_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3t4ZOmg5sTrx"
   },
   "source": [
    "Alright! The important configuration parameters are usually the number of layers `config.encoder_num_layers` and `config.decoder_num_layers`, the model's hidden size: `config.d_model`, the number of attention heads `config.encoder_attention_heads` and `config.decoder_attention_heads` and the vocabulary size `config.vocab_size`.\n",
    "\n",
    "Let's create 4 configurations different from the baseline and see how they compare in terms of peak memory consumption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qA0d1RvGYAEE"
   },
   "outputs": [],
   "source": [
    "config_baseline = BartConfig.from_pretrained(\"facebook/bart-large-mnli\")\n",
    "config_768_hidden = BartConfig.from_pretrained(\"facebook/bart-large-mnli\", d_model=768)\n",
    "config_8_heads = BartConfig.from_pretrained(\"facebook/bart-large-mnli\", decoder_attention_heads=8, encoder_attention_heads=8)\n",
    "config_10000_vocab = BartConfig.from_pretrained(\"facebook/bart-large-mnli\", vocab_size=10000)\n",
    "config_8_layers = BartConfig.from_pretrained(\"facebook/bart-large-mnli\", encoder_layers=8, decoder_layers=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RhefJji1rU07"
   },
   "source": [
    "Cool, now we can benchmark these configs against the baseline config. This time, instead of using the benchmarking script we will directly use the `PyTorchBenchmark` class. The class expects the argument `args` which has to be of type `PyTorchBenchmarkArguments` and optionally a list of configs.\n",
    "\n",
    "First, we define the `args` and give the different configurations appropriate model names. The model names must be in the same order as the configs that are directly passed to `PyTorchBenchMark`.\n",
    "\n",
    "If no `configs` are provided to `PyTorchBenchmark`, it is assumed that the model names `[\"bart-base\", \"bart-768-hid\", \"bart-8-head\", \"bart-10000-voc\", \"bart-8-lay\"]` correspond to official model identifiers and their corresponding configs are loaded as was shown in the previous section.\n",
    "\n",
    "It is assumed that the model will be trained on half-precision, so we add the option `fp16=True` for the following benchmarks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 554
    },
    "colab_type": "code",
    "id": "Lv_WvM2jr79r",
    "outputId": "939dc355-036f-45ad-c996-e6cb136c7a59"
   },
   "outputs": [],
   "source": [
    "# define args\n",
    "args = PyTorchBenchmarkArguments(models=[\"bart-base\", \"bart-768-hid\", \"bart-8-head\", \"bart-10000-voc\", \"bart-8-lay\"], \n",
    "                                 no_speed=True,\n",
    "                                 no_inference=True,\n",
    "                                 training=True, \n",
    "                                 train_memory_csv_file=\"plots_pt/training_mem_fp16.csv\", \n",
    "                                 save_to_csv=True, \n",
    "                                 env_info_csv_file=\"plots_pt/env.csv\",\n",
    "                                 sequence_lengths=[64, 128, 256, 512],\n",
    "                                 batch_sizes=[8],\n",
    "                                 no_env_print=True,\n",
    "                                 fp16=True)  # let's train on fp16\n",
    "\n",
    "# create benchmark\n",
    "benchmark = PyTorchBenchmark(configs=[config_baseline, config_768_hidden, config_8_heads, config_10000_vocab, config_8_layers], args=args)\n",
    "\n",
    "# run benchmark\n",
    "result = benchmark.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DJWs_tDjxzuO"
   },
   "source": [
    "Nice, let's plot the results again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 534
    },
    "colab_type": "code",
    "id": "0r-r-R1lxEr0",
    "outputId": "5dbeb7f7-c996-4db2-a560-735354a5b76f"
   },
   "outputs": [],
   "source": [
    "# plot graph and save as image\n",
    "!python plot_csv_file.py --csv_file plots_pt/training_mem_fp16.csv --figure_png_file=plots_pt/training_mem_fp16.png --no_log_scale\n",
    "\n",
    "# show image\n",
    "from IPython.display import Image\n",
    "Image('plots_pt/training_mem_fp16.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5xTuRPBCx-dw"
   },
   "source": [
    "As expected the model of the baseline config requires the most memory. \n",
    "\n",
    "It is interesting to see that the \"bart-8-head\" model initially requires more memory than `bart-10000-voc`, but then clearly outperforms `bart-10000-voc` at an input length of 512. \n",
    "Less surprising is that the \"bart-8-lay\" is by far the most memory-efficient model when reminding oneself that during the forward pass every layer has to store its activations for the backward pass.\n",
    "\n",
    "Alright, given the data above, let's say we narrow our candidates down to only the \"bart-8-head\" and \"bart-8-lay\" models. \n",
    " \n",
    "Let's compare these models again on training time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 269
    },
    "colab_type": "code",
    "id": "c9xSoCUZ0Hlz",
    "outputId": "7054af8a-3050-4aca-f503-e229ed365cb0"
   },
   "outputs": [],
   "source": [
    "# define args\n",
    "args = PyTorchBenchmarkArguments(models=[\"bart-8-head\", \"bart-8-lay\"], \n",
    "                                 no_inference=True,\n",
    "                                 training=True,\n",
    "                                 no_memory=True,\n",
    "                                 train_time_csv_file=\"plots_pt/training_speed_fp16.csv\", \n",
    "                                 save_to_csv=True, \n",
    "                                 env_info_csv_file=\"plots_pt/env.csv\",\n",
    "                                 sequence_lengths=[32, 128, 512],\n",
    "                                 batch_sizes=[8],\n",
    "                                 no_env_print=True,\n",
    "                                 repeat=1, # to make speed measurement faster but less accurate\n",
    "                                 no_multi_process=True,  # google colab has problems with multi processing\n",
    "                                 fp16=True\n",
    "                                 )\n",
    "\n",
    "# create benchmark\n",
    "benchmark = PyTorchBenchmark(configs=[config_8_heads, config_8_layers], args=args)\n",
    "\n",
    "# run benchmark\n",
    "result = benchmark.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UseFqLiuRQuX"
   },
   "source": [
    "The option `no_multi_process` disabled multi-processing here. This option should in general only be used for testing or debugging. Enabling multi-processing is crucial to ensure accurate memory consumption measurement, but is less important when only measuring speed. The main reason it is disabled here is that google colab sometimes raises \"CUDA initialization\" due to the notebook's environment. \n",
    "This problem does not arise when running benchmarks outside of a notebook.\n",
    "\n",
    "Alright, let's plot the last speed results as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 534
    },
    "colab_type": "code",
    "id": "8c6fjmWLU0Rx",
    "outputId": "8a4b4db7-abed-47c4-da61-c3b1ccae66f1"
   },
   "outputs": [],
   "source": [
    "# plot graph and save as image\n",
    "!python plot_csv_file.py --csv_file plots_pt/training_speed_fp16.csv --figure_png_file=plots_pt/training_speed_fp16.png --no_log_scale --is_time\n",
    "\n",
    "# show image\n",
    "from IPython.display import Image\n",
    "Image('plots_pt/training_speed_fp16.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b6T7I4lnVCpk"
   },
   "source": [
    "Unsurprisingly, \"bart-8-lay\" is faster than \"bart-8-head\" by a factor of ca. 1.3. It might very well be that reducing the layers by a factor of 2 leads to much more performance degradation than reducing the number of heads by a factor of 2.\n",
    "For more information on computational efficient Bart models, check out the new *distilbart* model [here](https://huggingface.co/models?search=distilbart)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "S4cG0NwfNugm"
   },
   "source": [
    "Alright, that's it! Now you should be able to benchmark your favorite models on your favorite configurations. \n",
    "\n",
    "Feel free to share your results with the community [here](https://github.com/huggingface/transformers/blob/master/examples/benchmarking/README.md) or by tweeting us https://twitter.com/HuggingFace 🤗."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Benchamrks",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
